import numpy as np
from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
from rouge import Rouge
from transformers import AutoTokenizer, EvalPrediction


def similarity_metrics(eval_pred:EvalPrediction, tokenizer:AutoTokenizer):
    predictions, labels = eval_pred
    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
    # Replace -100 in the labels as we can't decode them.
    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
 
    
    # 字符级别
    decoded_preds = [" ".join((pred.replace(" ", ""))) for pred in decoded_preds]
    decoded_labels = [" ".join((label.replace(" ", ""))) for label in decoded_labels]
    # 词级别，分词
    # decoded_preds = [" ".join(jieba.cut(pred.replace(" ", ""))) for pred in decoded_preds]
    # decoded_labels = [" ".join(jieba.cut(label.replace(" ", ""))) for label in decoded_labels]
    rouge = Rouge()
    labels_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in labels]
 
 
    total = 0
 
    rouge_1, rouge_2, rouge_l, bleu = 0, 0, 0, 0
    for decoded_label, decoded_pred in zip(decoded_labels, decoded_preds):
        total += 1
        scores = rouge.get_scores(hyps=decoded_pred, refs=decoded_label)
        rouge_1 += scores[0]['rouge-1']['f']
        rouge_2 += scores[0]['rouge-2']['f']
        rouge_l += scores[0]['rouge-l']['f']
        bleu += sentence_bleu(
            references=[decoded_label.split(' ')],
            hypothesis=decoded_pred.split(' '),
            smoothing_function=SmoothingFunction().method1
        )
    bleu /= len(decoded_labels)
    rouge_1 /= total
    rouge_2 /= total
    rouge_l /= total
    result = {'rouge-1': rouge_1, 'rouge-2': rouge_2, 'rouge-l': rouge_l}
    print(result)
    # 测试平均与分别计算是否一致
    result2 = rouge.get_scores(decoded_preds, decoded_labels, avg=True)
    print(result2)
    print(bleu)
    # result = {'rouge-1': result['rouge-1']['f'], 'rouge-2': result['rouge-2']['f'], 'rouge-l': result['rouge-l']['f']}
 
    result = {key: value * 100 for key, value in result.items()}
    result["gen_len"] = np.mean(labels_lens)
    result["bleu"] = bleu * 100
    return result